# K = 3 and d = 2, m = 8, n = 6, C^(A) = 4, C^(B) = 2
import numpy as np

# k = 1, d = 1
A_1_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_1_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 1, d = 2
A_1_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_1_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 2, d = 1
A_2_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_2_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 2, d = 2
A_2_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_2_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 3, d = 1
A_3_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_3_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# k = 3, d = 2
A_3_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (8, 4))
B_3_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (6, 2))

# x
X = np.random.uniform(0.0, 1.0, (8,8))


def normal_algo_1():
	AB_1_1 = np.kron(A_1_1, B_1_1)
	AB_1_2 = np.kron(A_1_2, B_1_2)
	AB_2_1 = np.kron(A_2_1, B_2_1)
	AB_2_2 = np.kron(A_2_2, B_2_2)
	AB_3_1 = np.kron(A_3_1, B_3_1)
	AB_3_2 = np.kron(A_3_2, B_3_2)

	W_1 = np.concatenate([AB_1_1, AB_2_1, AB_3_1], axis = -1)
	W_2 = np.concatenate([AB_1_2, AB_2_2, AB_3_2], axis = -1)

	# W_1 to (m_1 * n_1*C)
	W_1 = np.reshape(np.reshape(W_1, (8, 6, -1)), (8, -1))

	# Y is (m_2 * n_1*C)
	Y = np.matmul(X.T, W_1)

	# Y to (n_1 * m_2*C)
	Y = np.reshape(Y, (8, 6, -1))
	Y = np.transpose(Y, (1, 0, 2))
	Y = np.reshape(Y, (6, -1))

	# W_2 to (m_2*C * n_2)
	W_2 = np.reshape(W_2, (8, 6, -1))
	W_2 = np.transpose(W_2, (0, 2, 1))
	W_2 = np.reshape(W_2, (-1, 6))

	# Y is (n_1 * n_2)
	Y = np.matmul(Y, W_2)
	return Y


def divided_algo_1():
	def single_algo_1(I, W_1, W_2):
		# W_1 to (m_1 * n_1*C)
		W_1 = np.reshape(np.reshape(W_1, (8, 6, -1)), (8, -1))

		# O is (m_2 * n_1*C)
		O = np.matmul(I.T, W_1)

		# O to (n_1 * m_2*C)
		O = np.reshape(O, (8, 6, -1))
		O = np.transpose(O, (1, 0, 2))
		O = np.reshape(O, (6, -1))

		# W_2 to (m_2*C * n_2)
		W_2 = np.reshape(W_2, (8, 6, -1))
		W_2 = np.transpose(W_2, (0, 2, 1))
		W_2 = np.reshape(W_2, (-1, 6))

		# O is (n_1 * n_2)
		O = np.matmul(O, W_2)
		return O

	AB_1_1 = np.kron(A_1_1, B_1_1)
	AB_1_2 = np.kron(A_1_2, B_1_2)
	AB_2_1 = np.kron(A_2_1, B_2_1)
	AB_2_2 = np.kron(A_2_2, B_2_2)
	AB_3_1 = np.kron(A_3_1, B_3_1)
	AB_3_2 = np.kron(A_3_2, B_3_2)

	X_1 = single_algo_1(X, AB_1_1, AB_1_2)
	X_2 = single_algo_1(X, AB_2_1, AB_2_2)
	X_3 = single_algo_1(X, AB_3_1, AB_3_2)

	Y = X_1 + X_2 + X_3
	return Y


def divided_algo_2():
	def single_algo_2(I, A_1, B_1, A_2, B_2):
		# O is (m_2 * C_A)
		O = np.matmul(I.T, A_1)

		# O is (m_2*n_1 * C_A*C_B)
		O = np.kron(O, B_1)

		# O to (n_1*C_B * m_2*C_A)
		O = np.reshape(O, (8, 6, 4, 2))
		O = np.transpose(O, (1, 3, 0, 2))
		O = np.reshape(O, (6*2, 8*4))

		# O is (n_1 * C_B)
		O = np.matmul(O, np.reshape(A_2, (-1)))
		O = np.reshape(O, (6, 2))

		# O is (n_1 * n_2)
		O = np.matmul(O, B_2.T)
		return O

	X_1 = single_algo_2(X, A_1_1, B_1_1, A_1_2, B_1_2)
	X_2 = single_algo_2(X, A_2_1, B_2_1, A_2_2, B_2_2)
	X_3 = single_algo_2(X, A_3_1, B_3_1, A_3_2, B_3_2)

	Y = X_1 + X_2 + X_3
	return Y


if __name__ == '__main__':
	Y_1 = normal_algo_1()
	Y_2 = divided_algo_1()
	Y_3 = divided_algo_2()

	print('Normal is:')
	print(Y_1)
	print('\n')

	print('Divided 1 is:')
	print(Y_2)
	print('\n')

	print('Divided 2 is:')
	print(Y_3)
	print('\n')
